"""
Agent端文件接收服务器
专门用于接收控制端发送的文件，使用IP白名单保护
"""
import asyncio
import os
import uuid
from pathlib import Path
from typing import Optional, Set
from datetime import datetime, timedelta

from aiohttp import web, ClientError
from loguru import logger


class FileReceiverServer:
    """Agent端文件接收服务器"""
    
    def __init__(self, host: str = '0.0.0.0', port: int = 9002, files_dir: str = None, allowed_ips: Set[str] = None):
        self.host = host
        self.port = port
        self.files_dir = Path(files_dir or os.path.join(os.getcwd(), 'files'))
        self.allowed_ips = allowed_ips or set()
        
        # 创建文件目录结构
        self.control_uploads_dir = self.files_dir / 'control_uploads'  # 控制端发送的文件
        self.chat_files_dir = self.files_dir / 'chat_files'  # 闲鱼聊天接收的二进制文件
        self.control_uploads_dir.mkdir(parents=True, exist_ok=True)
        self.chat_files_dir.mkdir(parents=True, exist_ok=True)
        
        # 文件存储：file_id -> (file_path, created_at, file_type)
        self.files: dict = {}
        
        # aiohttp应用
        self.app = web.Application()
        self._setup_routes()
        self._runner = None
        self._site = None
    
    def update_allowed_ips(self, ips: Set[str]):
        """更新允许的IP列表（从TCP连接的控制端IP）"""
        self.allowed_ips = ips
        logger.info(f"✅ 文件接收服务器IP白名单已更新: {list(ips)}")
    
    def _check_ip_whitelist(self, request: web.Request) -> bool:
        """检查IP是否在白名单中"""
        # 获取客户端IP
        client_ip = request.remote
        
        # 如果白名单为空，拒绝所有请求（安全默认）
        if not self.allowed_ips:
            logger.warning(f"❌ 文件接收服务器IP白名单为空，拒绝请求: {client_ip}")
            return False
        
        # 检查IP是否在白名单中
        if client_ip in self.allowed_ips:
            return True
        
        logger.warning(f"❌ 文件接收服务器拒绝未授权IP: {client_ip}, 白名单: {list(self.allowed_ips)}")
        return False
    
    def _setup_routes(self):
        """设置路由"""
        
        async def upload_file(request: web.Request):
            """接收文件（控制端调用）"""
            # ✅ IP白名单验证
            if not self._check_ip_whitelist(request):
                return web.Response(
                    status=403,
                    text="Forbidden: IP not in whitelist"
                )
            
            try:
                # 读取文件
                reader = await request.multipart()
                field = await reader.next()
                
                if not field or field.name != 'file':
                    return web.Response(status=400, text="Missing file field")
                
                file_content = await field.read()
                file_name = field.filename or 'uploaded_file'
                
                # ✅ 文件大小限制（50MB）
                MAX_FILE_SIZE = 50 * 1024 * 1024
                if len(file_content) > MAX_FILE_SIZE:
                    return web.Response(
                        status=400,
                        text=f"File too large (max {MAX_FILE_SIZE // 1024 // 1024}MB)"
                    )
                
                # ✅ 文件类型验证（只允许图片）
                allowed_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.webp', '.bmp'}
                file_ext = Path(file_name).suffix.lower()
                if file_ext not in allowed_extensions:
                    return web.Response(
                        status=400,
                        text=f"File type not allowed. Allowed: {', '.join(allowed_extensions)}"
                    )
                
                # ✅ 生成唯一文件ID
                file_id = str(uuid.uuid4())
                
                # ✅ 保存文件到control_uploads目录
                file_path = self.control_uploads_dir / f"{file_id}_{file_name}"
                file_path.write_bytes(file_content)
                
                # 记录文件信息
                self.files[file_id] = {
                    'file_path': file_path,
                    'file_name': file_name,
                    'file_size': len(file_content),
                    'file_type': 'control_upload',
                    'created_at': datetime.now()
                }
                
                logger.info(f"✅ 文件已接收: file_id={file_id}, name={file_name}, size={len(file_content)}字节, from={request.remote}")
                
                return web.json_response({
                    'file_id': file_id,
                    'file_name': file_name,
                    'file_size': len(file_content),
                    'file_path': str(file_path)
                })
                
            except Exception as e:
                logger.error(f"接收文件失败: {e}", exc_info=True)
                return web.Response(status=500, text=str(e))
        
        async def get_file(request: web.Request):
            """获取文件（控制端或Agent内部调用）"""
            file_id = request.match_info.get('file_id')
            
            if not file_id or file_id not in self.files:
                return web.Response(status=404, text="File not found")
            
            file_info = self.files[file_id]
            file_path = file_info['file_path']
            
            if not file_path.exists():
                del self.files[file_id]
                return web.Response(status=404, text="File not found")
            
            # ✅ 返回文件
            return web.Response(
                body=file_path.read_bytes(),
                headers={
                    'Content-Type': 'application/octet-stream',
                    'Content-Disposition': f'attachment; filename="{file_info["file_name"]}"'
                }
            )
        
        async def delete_file(request: web.Request):
            """删除文件（清理）"""
            file_id = request.match_info.get('file_id')
            
            if not file_id or file_id not in self.files:
                return web.json_response({'success': False, 'msg': 'File not found'})
            
            file_info = self.files[file_id]
            file_path = file_info['file_path']
            
            # 删除文件
            if file_path.exists():
                try:
                    file_path.unlink()
                except Exception as e:
                    logger.warning(f"删除文件失败: {e}")
            
            del self.files[file_id]
            logger.info(f"🗑️ 文件已删除: file_id={file_id}")
            
            return web.json_response({'success': True, 'msg': 'File deleted'})
        
        async def get_status(request: web.Request):
            """获取服务器状态"""
            return web.json_response({
                'files_count': len(self.files),
                'allowed_ips': list(self.allowed_ips),
                'files_dir': str(self.files_dir),
                'host': self.host,
                'port': self.port
            })
        
        # 注册路由
        self.app.router.add_post('/upload', upload_file)
        self.app.router.add_get('/file/{file_id}', get_file)
        self.app.router.add_delete('/file/{file_id}', delete_file)
        self.app.router.add_get('/status', get_status)
    
    async def start_server(self):
        """启动文件接收服务器"""
        self._runner = web.AppRunner(self.app)
        await self._runner.setup()
        self._site = web.TCPSite(self._runner, self.host, self.port)
        await self._site.start()
        logger.info(f"📁 Agent文件接收服务器已启动: {self.host}:{self.port}, 允许IP: {list(self.allowed_ips)}")
    
    async def stop_server(self):
        """停止文件接收服务器"""
        if self._site:
            await self._site.stop()
        if self._runner:
            await self._runner.cleanup()
        logger.info("Agent文件接收服务器已停止")
    
    def get_file_path(self, file_id: str) -> Optional[Path]:
        """获取文件路径（如果存在）"""
        if file_id in self.files:
            file_info = self.files[file_id]
            file_path = file_info['file_path']
            if file_path.exists():
                return file_path
        return None
    
    def save_chat_file(self, file_content: bytes, file_name: str, file_type: str = 'chat') -> str:
        """保存闲鱼聊天接收的二进制文件"""
        file_id = str(uuid.uuid4())
        file_path = self.chat_files_dir / f"{file_id}_{file_name}"
        file_path.write_bytes(file_content)
        
        self.files[file_id] = {
            'file_path': file_path,
            'file_name': file_name,
            'file_size': len(file_content),
            'file_type': file_type,
            'created_at': datetime.now()
        }
        
        logger.info(f"✅ 聊天文件已保存: file_id={file_id}, name={file_name}, size={len(file_content)}字节")
        return file_id

